import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from torch.utils.data.sampler import BatchSampler, SubsetRandomSampler
from torch.distributions import Categorical
from masking_categorical import CategoricalMasked
import config
from object import delay
import random
import math
import os

import pandas as pd
import dcn as dcn
from sklearn.preprocessing import LabelEncoder, MinMaxScaler
from deepctr_torch.models import xDeepFM
from deepctr_torch.inputs import SparseFeat, DenseFeat, get_feature_names

gpus = [0]
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

# Trick 8: orthogonal initialization
def orthogonal_init(layer, gain=1.0):
    nn.init.orthogonal_(layer.weight, gain=gain)
    nn.init.constant_(layer.bias, 0)
    
import torch
import torch.nn as nn
import torch.nn.functional as F

class SimpleDiffusionModel(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, T=1000):
        super(SimpleDiffusionModel, self).__init__()
        self.T = T  # 总的扩散步骤数
        self.beta = torch.linspace(1e-4, 0.02, T).to(device)  # 线性beta调度
        self.alpha = 1.0 - self.beta
        self.alpha_bar = torch.cumprod(self.alpha, dim=0)
        
        # 输入维度包括 x 和 t
        self.model = nn.Sequential(
            nn.Linear(input_dim + 1, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, output_dim)
        )
    
    def forward(self, x, t):
        # 确保 t 是 [batch_size, 1]
        t = t.unsqueeze(1).float()
        out = torch.cat([x, t], dim=1)
        output = self.model(out)
        
        # 检查输出是否有 NaN
        if torch.isnan(output).any():
            print("Warning: Diffusion model output contains NaN values.")
            # 替换 NaN 为零
            output = torch.where(torch.isnan(output), torch.zeros_like(output), output)
        
        return output
    
    def sample(self, x_start):
        # 从噪声中逐步生成样本
        x = x_start
        for t in reversed(range(self.T)):
            beta_t = self.beta[t]
            alpha_t = self.alpha[t]
            alpha_bar_t = self.alpha_bar[t]
            # 预测噪声
            eps = self.forward(x, torch.tensor([t]).to(device))
            # 计算均值
            mean = (x - beta_t * eps) / torch.sqrt(alpha_t)
            # 添加噪声
            if t > 0:
                noise = torch.randn_like(x)
                x = mean + torch.sqrt(beta_t) * noise
            else:
                x = mean
        return x

# Actor类集成扩散模型


class Actor(nn.Module):
    def __init__(self, args):
        super(Actor, self).__init__()
        
        self.fc_downloadtime = nn.Linear(9, 512)
        self.fc_alg_shu = nn.Linear(731, 512)
        self.fc2 = nn.Linear(512, 128)
        self.fc3 = nn.Linear(128, 128)
        self.fc4 = nn.Linear(128, args.action_dim)
        self.linear_layer = nn.Linear(32, 512)
        self.linear_layer_2 = nn.Linear(10240, 10)
        
        self.activate_func = [nn.ReLU(), nn.Tanh()][args.use_tanh]
        self.deepdcn = None
        self.embedding_size = 64

        # 定义扩散模型
        self.diffusion_model = SimpleDiffusionModel(
            input_dim=args.action_dim,  # 10
            hidden_dim=256,
            output_dim=args.action_dim,  # 10
            T=1000
        ).to(device)

        # 参数初始化
        if args.use_orthogonal_init:
            print("------use_orthogonal_init------")
            orthogonal_init(self.fc_downloadtime)
            orthogonal_init(self.fc_alg_shu)
            orthogonal_init(self.fc2)
            orthogonal_init(self.fc3)
            orthogonal_init(self.fc4, gain=0.01)

        # 定义特征
        self.dense_features = ['C' + str(i) for i in range(1, 13)]
        self.sparse_features = ['I' + str(i) for i in range(1, 720)]

    def forward(self, s):
        # 现有的前向传播逻辑
        s = s.view(-1)
        download_finish_time = s[:9].reshape(1, 9)
        alg_shu = s[9:9 + 731 * 10].reshape(10, 731)
        tmp_layer01 = s[9 + 731 * 10:9 + 731 * 10 + 719 * 9].reshape(9, 719)

        downloadtime_input = torch.tensor(download_finish_time).view(1, -1).to(device)
        downloadtime_input = downloadtime_input.to(self.fc_downloadtime.weight.dtype)
        layer01_input = torch.tensor(tmp_layer01, dtype=torch.float).to(device)
        alg_shu_input = torch.tensor(alg_shu, dtype=torch.float).to(device)
        s_ = np.resize(alg_shu_input.cpu().detach().numpy(), (config.EDGE_NODE_NUM + 1, 731))
        s_ = torch.tensor(s_).to(device).float()

        self.s_df = pd.DataFrame(s_.cpu().numpy(), columns=self.dense_features + self.sparse_features)
        feat_sizes = {}
        feat_sizes_dense = {feat: 1 for feat in self.dense_features}
        feat_sizes_sparse = {feat: len(self.s_df[feat].unique()) for feat in self.sparse_features}
        feat_sizes.update(feat_sizes_dense)
        feat_sizes.update(feat_sizes_sparse)

        for feat in self.sparse_features:
            lbe = LabelEncoder()
            self.s_df[feat] = lbe.fit_transform(self.s_df[feat])

        mms = MinMaxScaler(feature_range=(0, 1))
        self.s_df[self.dense_features] = mms.fit_transform(self.s_df[self.dense_features])

        self.fixlen_feature_columns = [(feat, 'sparse') for feat in self.sparse_features] + [(feat, 'dense') for feat in self.dense_features]
        self.dnn_feature_columns = self.fixlen_feature_columns
        self.linear_feature_columns = self.fixlen_feature_columns

        if self.deepdcn is None:
            self.deepdcn = dcn.DCN(feat_sizes, self.embedding_size, self.linear_feature_columns, self.dnn_feature_columns).to(device)
        
        downloadtime_output = self.activate_func(self.fc_downloadtime(downloadtime_input))
        layer01_output = self.deepdcn(layer01_input)
        alg_shu_fc = self.activate_func(self.fc_alg_shu(s_).squeeze(1))

        downloadtime_output = downloadtime_output.view(1, 512)
        layer01_output = layer01_output.view(9, 32)
        alg_shu_fc = alg_shu_fc.view(10, 512)

        layer01_output = layer01_output.to(device)
        layer01_output = self.linear_layer(layer01_output)

        combined_output = torch.cat((downloadtime_output, layer01_output, alg_shu_fc), dim=0)
        a_prob_initial = combined_output.view(1, -1)
        a_prob_initial = torch.softmax(self.linear_layer_2(a_prob_initial), dim=1)  # [1,10]

        # 使用扩散模型生成动作分布
        t = torch.randint(0, self.diffusion_model.T, (a_prob_initial.size(0),)).to(device)  # [1]
        a_prob_denoised = self.diffusion_model(a_prob_initial, t)  # [1,10]

        # 检查扩散模型的输出
        if torch.isnan(a_prob_denoised).any():
            print("Warning: Diffusion model output contains NaN values.")
            a_prob_denoised = torch.where(torch.isnan(a_prob_denoised), torch.zeros_like(a_prob_denoised), a_prob_denoised)

        a_prob = torch.softmax(a_prob_denoised, dim=1)

        return a_prob




class Critic(nn.Module):
    def __init__(self, args):
        super(Critic, self).__init__()
        self.fc1 = nn.Linear(13790, args.hidden_width)
        self.fc2 = nn.Linear(args.hidden_width, args.hidden_width)
        self.fc3 = nn.Linear(args.hidden_width, 1)
        self.activate_func = [nn.ReLU(), nn.Tanh()][args.use_tanh]  

        if args.use_orthogonal_init:
            print("------use_orthogonal_init------")
            orthogonal_init(self.fc1)
            orthogonal_init(self.fc2)
            orthogonal_init(self.fc3)

    def forward(self, s):
        s = torch.flatten(s, start_dim=1)
        
        s = self.activate_func(self.fc1(s))
        s = self.activate_func(self.fc2(s))
        v_s = self.fc3(s)
        return v_s

class RunningMeanStd:
    # Dynamically calculate mean and std
    def __init__(self, shape):  # shape:the dimension of input data
        self.n = 0
        self.mean = np.zeros(shape)
        self.S = np.zeros(shape)
        self.std = np.sqrt(self.S)

    def update(self, x):
        x = np.array(x)
        self.n += 1
        if self.n == 1:
            self.mean = x
            self.std = x
        else:
            old_mean = self.mean.copy()
            self.mean = old_mean + (x - old_mean) / self.n
            self.S = self.S + (x - old_mean) * (x - self.mean)
            self.std = np.sqrt(self.S / self.n)


class RewardScaling:
    def __init__(self, shape, gamma):
        self.shape = shape  # reward shape=1
        self.gamma = gamma  # discount factor
        self.running_ms = RunningMeanStd(shape=self.shape)
        self.R = np.zeros(self.shape)

    def __call__(self, x):
        self.R = self.gamma * self.R + x
        self.running_ms.update(self.R)
        x = x / (self.running_ms.std + 1e-8)  # Only divided std
        return x

    def reset(self):  # When an episode is done,we should reset 'self.R'
        self.R = np.zeros(self.shape)


class PPO_discrete:
    def __init__(self, args):
        self.batch_size = args.batch_size
        self.batch_size_pre = args.batch_size_pre
        self.mini_batch_size = args.mini_batch_size
        self.max_train_steps = args.max_train_steps
        self.lr = args.lr
        self.lr_a = args.lr_a  # Learning rate of actor
        self.lr_c = args.lr_c  # Learning rate of critic
        self.gamma = args.gamma  # Discount factor
        self.lamda = args.lamda  # GAE parameter
        self.epsilon = args.epsilon  # PPO clip parameter
        self.K_epochs = args.K_epochs  # PPO parameter
        self.entropy_coef = args.entropy_coef  # Entropy coefficient
        self.set_adam_eps = args.set_adam_eps
        self.use_grad_clip = args.use_grad_clip
        self.use_lr_decay = args.use_lr_decay
        self.use_adv_norm = args.use_adv_norm

        self.actor = Actor(args).to(device)
        self.critic = Critic(args).to(device)
        self.rewardscaling = RewardScaling(1, self.gamma)

        if self.set_adam_eps:  
            self.optimizer_actor = torch.optim.Adam(self.actor.parameters(), lr=self.lr_a, eps=1e-5)
            self.optimizer_critic = torch.optim.Adam(self.critic.parameters(), lr=self.lr_c, eps=1e-5)
        else:
            self.optimizer_actor = torch.optim.Adam(self.actor.parameters(), lr=self.lr_a)
            self.optimizer_critic = torch.optim.Adam(self.critic.parameters(), lr=self.lr_c)
        
        # 引入扩散模型的优化器
        self.optimizer_diffusion = torch.optim.Adam(
            self.actor.diffusion_model.parameters(),
            lr=args.lr  # 可以单独设置学习率
        )
        
        # 扩散模型的损失函数
        self.diffusion_loss_fn = nn.MSELoss()
    

    def choose_action(self, s, uid, obs):
        available_actions = obs['next_can']
        ttt = []
        for id, item in enumerate(available_actions):
            if item == 1:
                ttt.append(id)
        # print('(choose_action)available_actions: ',ttt)
    
        s_tensor = torch.unsqueeze(torch.tensor(s, dtype=torch.float), 0).to(device)
    
        with torch.no_grad():
            available_actions_tensor = torch.tensor(available_actions, dtype=torch.bool).to(device)
            a_prob_initial = self.actor(s_tensor)  # [1,10]
            
            # 使用扩散模型生成动作分布
            t = torch.randint(0, self.actor.diffusion_model.T, (a_prob_initial.size(0),)).to(device)  # [1]
            a_prob_denoised = self.actor.diffusion_model(a_prob_initial, t)  # [1,10]
            
            # 检查扩散模型的输出
            if torch.isnan(a_prob_denoised).any():
                print("Warning: Diffusion model output contains NaN values.")
                a_prob_denoised = torch.where(torch.isnan(a_prob_denoised), torch.zeros_like(a_prob_denoised), a_prob_denoised)
            
            a_prob = torch.softmax(a_prob_denoised, dim=1)  # [1,10]
    
            a_prob_refined = a_prob.squeeze(0)  # [10]
            a_prob_refined = a_prob_refined * available_actions_tensor.float()  # [10]
            a_prob_sum = a_prob_refined.sum()
    
            if a_prob_sum.item() > 0:
                a_prob_refined = a_prob_refined / (a_prob_sum + 1e-8)
            else:
                print("Warning: All actions are masked. Assigning uniform probabilities over all actions.")
                a_prob_refined = torch.ones_like(a_prob_refined) / a_prob_refined.size(-1)  # [10]
    
            # 检查 NaN
            if torch.isnan(a_prob_refined).any():
                print("Warning: a_prob_refined contains NaN values.")
                a_prob_refined = torch.where(torch.isnan(a_prob_refined), torch.zeros_like(a_prob_refined), a_prob_refined)
    
            # 使用 argmax 获取动作
            a = torch.argmax(a_prob_refined)
    
            dist = Categorical(probs=a_prob_refined)
    
        if len(ttt) != 0:
            action_retry_count = 0
            usr_locx = [j for i, j, k in obs['each_use_loc'] if i == uid][0]
            usr_locy = [k for i, j, k in obs['each_use_loc'] if i == uid][0]
    
            while self.satisfied_constrain(a.item(), uid, obs) == False:  # not satisfy constraints
                a = dist.sample()  # arbitrary prob distribution
                action_retry_count += 1
                if action_retry_count >= config.EDGE_NODE_NUM or usr_locx == np.inf or usr_locy == np.inf:
                    # raise EnvironmentError
                    a = config.EDGE_NODE_NUM  # to cloud
                    a = torch.tensor([a]).to(device)
                    break
        else:
            a = config.EDGE_NODE_NUM
            a = torch.tensor([a]).to(device)
    
        with torch.no_grad():
            a_logprob = dist.log_prob(a)
    
        a = a.item()
        a_logprob = a_logprob.item()
        a_prob_initial = a_prob_initial.squeeze(0).cpu().numpy()
        return a, a_logprob,a_prob_initial



    
    def choose_action_2(self, s,uid,obs,target):
        available_actions = obs['next_can']
        ttt = []
        for id,item in enumerate(available_actions):
            if item == 1:
                ttt.append(id)
        # print('(choose_action)available_actions: ',ttt)

        s = torch.unsqueeze(torch.tensor(s, dtype=torch.float), 0).to(device)
        # exit(39)
        with torch.no_grad():
            available_actions = torch.tensor(available_actions,dtype=torch.bool)
            probs=self.actor(s)
            dist = CategoricalMasked(logits=probs,mask=available_actions)
            # a = dist.sample()
            # a = dist.argmax()
            
        a = target
        a = torch.tensor([a])

        with torch.no_grad():
            a = a.to(device)
            a_logprob = dist.log_prob(a)

        a = a.item()
        a_logprob = a_logprob.item()
        # return a.numpy()[0], a_logprob.numpy()[0]
        return a, a_logprob
    
    def satisfied_constrain(self, edge_id, uidre, obs=None):
        # original from baseline.py
        # print("each edge_id is: ", edge_id)
        one_ulyr = [j for i,j in obs['each_usrhas_lyer'] if i == uidre][0]
        if edge_id == config.EDGE_NODE_NUM:
            return True
        edge_con = [d for a,b,c,d,f in obs['each_edge_cpumemcondisk'] if a == edge_id][0]
        if config.node_max_container_number - edge_con -1 < 0:
            print('edge {} container num is {} > 5.'.format(edge_id,edge_con))
            return  False
        # edge_usrnum_limit
        if config.node_usrnum_limit - obs['each_edge_usrnum'][edge_id] -1 < 0:
            print('edge {} has user num {} > 50.'.format(edge_id,obs['each_edge_usrnum'][edge_id]))
            return  False
        # edge_storage_free - task_size - download_size
        tmp_downsize = 0
        for item in one_ulyr:
            if item not in obs['each_edgehas_lyr'][edge_id]:
                tmp_downsize += obs['all_layer_downsiz'][item]

        edge_disk = [f for a,b,c,d,f in obs['each_edge_cpumemcondisk'] if a == edge_id][0]

        usr_has_tsksiz = [k for i,j,k in obs['each_usrhas_task'] if i == uidre][0]

        # print("edge {a} has disk {b} - layer_download_size:{d}".format(a=edge_id, b=edge_disk, d=tmp_downsize))
        if edge_disk - usr_has_tsksiz - tmp_downsize <0:
            print('edge {a} has disk {b} - user_tsk_size:{c} - layer_download_size:{d} < 0.'.format(a=edge_id,b=edge_disk,c=usr_has_tsksiz,d=tmp_downsize))
            return  False

    # def update(self, replay_buffer, total_steps):
    #     s, a, a_logprob, r, s_, dw, done = replay_buffer.numpy_to_tensor()
    #     """
    #         Calculate the advantage using GAE
    #         'dw=True' means dead or win, there is no next state s'
    #         'done=True' represents the terminal of an episode(dead or win or reaching the max_episode_steps). When calculating the adv, if done=True, gae=0
    #     """
    #     a = a.to(device)
    #     s_ = s_.to(device)
    #     r = r.to(device)
    #     s = s.to(device)
    #     dw = dw.to(device)
    #     a_logprob = a_logprob.to(device)
    #     adv = []
    #     gae = 0
    #     with torch.no_grad():
    #         a = a.to(device)
    #         s_ = s_.to(device)
    #         r = r.to(device)
    #         s = s.to(device)
    #         dw = dw.to(device)
    #         a_logprob = a_logprob.to(device)

    #         vs = self.critic(s)
    #         vs_ = self.critic(s_)
    #         deltas = r + self.gamma * (1.0 - dw) * vs_ - vs
    #         for delta, d in zip(reversed(deltas.flatten().cpu().numpy()), reversed(done.flatten().cpu().numpy())):
    #             gae = delta + self.gamma * gae * (1.0 - d)
    #             adv.insert(0,gae)
    #         adv = torch.tensor(adv, dtype=torch.float).view(-1,1).to(device)
    #         v_target = adv + vs
    #         if self.use_adv_norm:
    #             adv = ((adv - adv.mean()) / (adv.std() + 1e-5))

    #     for _ in range(self.K_epochs):
    #         for index in BatchSampler(SubsetRandomSampler(range(self.batch_size)), self.mini_batch_size, False):
    #             dist_now = Categorical(probs=self.actor(s[index]))
    #             dist_entropy = dist_now.entropy().view(-1,1)
    #             a_logprob_now = dist_now.log_prob(a[index].squeeze()).view(-1,1).to(device)
    #             radios = torch.exp(a_logprob_now - a_logprob[index])

    #             surr1 = radios * adv[index]
    #             surr2 = torch.clamp(radios, 1 - self.epsilon, 1 + self.epsilon) * adv[index]
    #             actor_loss = -torch.min(surr1,surr2) - self.entropy_coef * dist_entropy

    #             # self.optimizer_actor.zero_grad()
    #             actor_loss = actor_loss.mean()
    #             self.optimizer_actor.zero_grad()
    #             actor_loss.backward()
    #             if self.use_grad_clip:
    #                 torch.nn.utils.clip_grad_norm_(self.actor.parameters(), 0.5)
    #             self.optimizer_actor.step()

    #             v_s = self.critic(s[index])
    #             critic_loss = F.mse_loss(v_target[index], v_s)

    #             self.optimizer_critic.zero_grad()
    #             critic_loss.backward()
    #             if self.use_grad_clip:
    #                 torch.nn.utils.clip_grad_norm_(self.critic.parameters(), 0.5)
    #             self.optimizer_critic.step()

    #     if self.use_lr_decay:
    #         self.lr_decay(total_steps)
        
    #     # 获取Replay Buffer中的动作分布
    #     actions = a  # 从Replay Buffer中获取动作
    #     # 将动作转换为one-hot编码
    #     target_action_distribution = F.one_hot(actions, num_classes=self.actor.fc4.out_features).float()
        
    #     # 使用扩散模型生成的动作分布
    #     a_prob_generated = self.actor.diffusion_model.sample(torch.randn_like(target_action_distribution))
        
    #     # 计算扩散模型的损失
    #     diffusion_loss = self.diffusion_loss_fn(a_prob_generated, target_action_distribution)
        
    #     # 反向传播扩散模型的损失
    #     self.optimizer_diffusion.zero_grad()
    #     diffusion_loss.backward()
    #     self.optimizer_diffusion.step()

    #     return actor_loss, critic_loss
    def update(self, replay_buffer, total_steps):
        # 从Replay Buffer中获取数据
        s, a, a_logprob, r, s_, dw, done, a_prob_initial = replay_buffer.numpy_to_tensor()
        a = a.to(device)
        s_ = s_.to(device)
        r = r.to(device)
        s = s.to(device)
        dw = dw.to(device)
        a_logprob = a_logprob.to(device)
        a_prob_initial = a_prob_initial.to(device)
    
        # 计算优势值（GAE）
        adv = []
        gae = 0
        with torch.no_grad():
            vs = self.critic(s)
            vs_ = self.critic(s_)
            deltas = r + self.gamma * (1.0 - dw) * vs_ - vs
            for delta, d in zip(reversed(deltas.flatten().cpu().numpy()), reversed(done.flatten().cpu().numpy())):
                gae = delta + self.gamma * gae * (1.0 - d)
                adv.insert(0, gae)
            adv = torch.tensor(adv, dtype=torch.float).view(-1, 1).to(device)
            v_target = adv + vs
            if self.use_adv_norm:
                adv = ((adv - adv.mean()) / (adv.std() + 1e-5))
    
        # PPO的多轮更新
        actor_loss_total = 0
        critic_loss_total = 0
        diffusion_loss_total = 0
        for _ in range(self.K_epochs):
            for index in BatchSampler(SubsetRandomSampler(range(self.batch_size)), self.mini_batch_size, False):
                # PPO Actor的损失
                dist_now = Categorical(probs=self.actor(s[index]))
                dist_entropy = dist_now.entropy().view(-1, 1)
                a_logprob_now = dist_now.log_prob(a[index].squeeze()).view(-1, 1).to(device)
                ratios = torch.exp(a_logprob_now - a_logprob[index])
    
                surr1 = ratios * adv[index]
                surr2 = torch.clamp(ratios, 1 - self.epsilon, 1 + self.epsilon) * adv[index]
                actor_loss = -torch.min(surr1, surr2) - self.entropy_coef * dist_entropy
    
                actor_loss = actor_loss.mean()
                self.optimizer_actor.zero_grad()
                actor_loss.backward()
                if self.use_grad_clip:
                    torch.nn.utils.clip_grad_norm_(self.actor.parameters(), 0.5)
                self.optimizer_actor.step()
    
                # PPO Critic的损失
                v_s = self.critic(s[index])
                critic_loss = F.mse_loss(v_target[index], v_s)
    
                self.optimizer_critic.zero_grad()
                critic_loss.backward()
                if self.use_grad_clip:
                    torch.nn.utils.clip_grad_norm_(self.critic.parameters(), 0.5)
                self.optimizer_critic.step()
    
                # 累计损失用于记录
                actor_loss_total += actor_loss.item()
                critic_loss_total += critic_loss.item()
    
        # 平均损失
        actor_loss_avg = actor_loss_total / (self.K_epochs * (self.batch_size // self.mini_batch_size))
        critic_loss_avg = critic_loss_total / (self.K_epochs * (self.batch_size // self.mini_batch_size))
    
        # 训练扩散模型
        # 使用存储的初始动作概率作为目标
        # 确保 a_prob_initial 不包含 NaN
        # 确保 a_prob_initial 不包含 NaN
        if torch.isnan(a_prob_initial).any():
            print("Warning: a_prob_initial contains NaN values. Replacing NaN with zeros.")
            a_prob_initial = torch.where(torch.isnan(a_prob_initial), torch.zeros_like(a_prob_initial), a_prob_initial)
        
        # 生成新的噪声时间步
        t = torch.randint(0, self.actor.diffusion_model.T, (a_prob_initial.size(0),)).to(device)
        
        # 计算扩散模型的输出
        a_prob_denoised = self.actor.diffusion_model(a_prob_initial, t)  # [batch_size, action_dim]
        
        # 计算扩散模型的损失
        diffusion_loss = F.mse_loss(a_prob_denoised, a_prob_initial)  # 示例损失函数
        
        # 确保损失不包含 NaN
        if torch.isnan(diffusion_loss):
            print("Warning: diffusion_loss is NaN. Setting loss to zero.")
            diffusion_loss = torch.tensor(0.0, device=device)
        
        self.optimizer_diffusion.zero_grad()
        diffusion_loss.backward()
        self.optimizer_diffusion.step()
        
        diffusion_loss_avg = diffusion_loss.item()

    
    def update1(self, replay_buffer, total_steps):
        s, a, a_logprob, r, s_, dw, done = replay_buffer.numpy_to_tensor()
        """
            Calculate the advantage using GAE
            'dw=True' means dead or win, there is no next state s'
            'done=True' represents the terminal of an episode(dead or win or reaching the max_episode_steps). When calculating the adv, if done=True, gae=0
        """
        a = a.to(device)
        s_ = s_.to(device)
        r = r.to(device)
        s = s.to(device)
        dw = dw.to(device)
        a_logprob = a_logprob.to(device)
        adv = []
        gae = 0
        with torch.no_grad():
            a = a.to(device)
            s_ = s_.to(device)
            r = r.to(device)
            s = s.to(device)
            dw = dw.to(device)
            a_logprob = a_logprob.to(device)

            vs = self.critic(s)
            vs_ = self.critic(s_)
            deltas = r + self.gamma * (1.0 - dw) * vs_ - vs
            for delta, d in zip(reversed(deltas.flatten().cpu().numpy()), reversed(done.flatten().cpu().numpy())):
                gae = delta + self.gamma * gae * (1.0 - d)
                adv.insert(0,gae)
            adv = torch.tensor(adv, dtype=torch.float).view(-1,1).to(device)
            v_target = adv + vs
            if self.use_adv_norm:
                adv = ((adv - adv.mean()) / (adv.std() + 1e-5))

        for _ in range(self.K_epochs):
            for index in BatchSampler(SubsetRandomSampler(range(self.batch_size_pre)), self.mini_batch_size, False):
                # print("s[index] is : ",s[index])
                # print(len(s[index]))
                # print(len(s))
                # print(len(s[index][0]))
                dist_now = Categorical(probs=self.actor(s[index]))
                dist_entropy = dist_now.entropy().view(-1,1)
                a_logprob_now = dist_now.log_prob(a[index].squeeze()).view(-1,1).to(device)
                radios = torch.exp(a_logprob_now - a_logprob[index])

                surr1 = radios * adv[index]
                surr2 = torch.clamp(radios, 1 - self.epsilon, 1 + self.epsilon) * adv[index]
                actor_loss = -torch.min(surr1,surr2) - self.entropy_coef * dist_entropy

                # self.optimizer_actor.zero_grad()
                actor_loss = actor_loss.mean()
                self.optimizer_actor.zero_grad()
                actor_loss.backward()
                if self.use_grad_clip:
                    torch.nn.utils.clip_grad_norm_(self.actor.parameters(), 0.5)
                self.optimizer_actor.step()

                v_s = self.critic(s[index])
                critic_loss = F.mse_loss(v_target[index], v_s)

                self.optimizer_critic.zero_grad()
                critic_loss.backward()
                if self.use_grad_clip:
                    torch.nn.utils.clip_grad_norm_(self.critic.parameters(), 0.5)
                self.optimizer_critic.step()

        if self.use_lr_decay:
            self.lr_decay(total_steps)

        return actor_loss, critic_loss

    def lr_decay(self, total_steps):
        
        if total_steps == 100:
            
            lr_a_now = 3e-4
            lr_c_now = 3e-4
            lr_now = 3e-4
            
        else:
            lr_a_now = self.lr_a * (1 - total_steps / self.max_train_steps)
            lr_c_now = self.lr_c * (1 - total_steps / self.max_train_steps)
            lr_now = self.lr * (1 - total_steps / self.max_train_steps)
        
        for p in self.optimizer_actor.param_groups:
            p['lr'] = lr_a_now
        for p in self.optimizer_critic.param_groups:
            p['lr'] = lr_c_now
        for p in self.optimizer_actor.param_groups:
            p['lr'] = lr_now

    def lr_decay1(self,total_steps):
        lr_a_now = self.lr_a * math.pow(0.95,total_steps/45)
        lr_c_now = self.lr_c * math.pow(0.95,total_steps/45)
        lr_now = self.lr * math.pow(0.98,total_steps/45)
        for p in self.optimizer_actor.param_groups:
            p['lr'] = lr_a_now
        for p in self.optimizer_critic.param_groups:
            p['lr'] = lr_c_now
        for p in self.optimizer_actor.param_groups:
            p['lr'] = lr_now

    def save_model(self, total_epochs,path):
        if not os.path.exists(path):
            open(path, 'a').close()
        checkpoint = {
            "net_actor": self.actor.state_dict(),
            "net_ctritic": self.critic.state_dict(),
            "optimizer": self.optimizer_actor.state_dict(),
            "optimizer": self.optimizer_critic.state_dict(),
            "epoch": total_epochs,
        }
        checkpoint["diffusion_model"] = self.actor.diffusion_model.state_dict()
        torch.save(checkpoint, path)

    def load_model(self, path):
        checkpoint = torch.load(path)
        
        # self.ac.load_state_dict(checkpoint["net"])
        self.actor.load_state_dict(checkpoint["net"])
        self.critic.load_state_dict(checkpoint["net"])
        self.optimizer_actor.load_state_dict(checkpoint["optimizer"])
        self.actor.diffusion_model.load_state_dict(checkpoint["diffusion_model"])
        epoch = checkpoint["epoch"]
        return epoch
